import os
import torch
import numpy as np
import cv2
import matplotlib
matplotlib.use('Agg')  # 设置非GUI后端
from matplotlib import pyplot as plt
from PIL import Image
from flask import Flask, request, render_template, jsonify
from torchvision import transforms
from models import get_modified_resnet18

app = Flask(__name__)
app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0  # 禁止缓存

# 加载模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = get_modified_resnet18().to(device)
model.load_state_dict(torch.load('best_model.pth', map_location=device))
model.eval()

# 预处理变换
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
])

# 类别名称（示例）
ID2CN = ['水仙花', '雪花莲', '雏菊', '款冬', '蒲公英', '黄花九轮草',
        '毛茛', '银莲花', '三色堇', '铃兰', '风铃草', '番红花',
        '鸢尾花', '虎皮百合', '郁金香', '棋盘花', '向日葵']

def generate_grad_cam(model, img_tensor):
    model.train()
    img_tensor.requires_grad = True

    # 捕获激活和梯度的钩子
    activations = None
    gradients = None

    def forward_hook(module, input, output):
        nonlocal activations
        activations = output.detach()

    def backward_hook(module, grad_input, grad_output):
        nonlocal gradients
        gradients = grad_output[0].detach()

    # 找到最后一个卷积层
    last_conv_layer = None
    for module in model.modules():
        if isinstance(module, torch.nn.Conv2d):
            last_conv_layer = module
    if last_conv_layer is None:
        raise ValueError("No Conv2d layer found")

    # 注册钩子
    forward_handle = last_conv_layer.register_forward_hook(forward_hook)
    backward_handle = last_conv_layer.register_backward_hook(backward_hook)

    # 前向传播
    output = model(img_tensor)
    pred_class = output.argmax(dim=1).item()
    target = output[0, pred_class]

    # 反向传播
    model.zero_grad()
    target.backward()

    # 移除钩子
    forward_handle.remove()
    backward_handle.remove()

    # 计算梯度均值
    pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])  # [out_channels]

    # 加权激活图
    activations = activations.squeeze(0)  # [out_channels, H, W]
    heatmap = torch.zeros(activations.shape[1:], dtype=torch.float32).to(activations.device)
    for i in range(activations.shape[0]):
        heatmap += pooled_gradients[i] * activations[i, :, :]

    # 后处理
    heatmap = torch.relu(heatmap).cpu().numpy()
    heatmap = (heatmap - np.min(heatmap)) / (np.max(heatmap) - np.min(heatmap) + 1e-8)
    heatmap = cv2.resize(heatmap, (224, 224))

    return heatmap, pred_class

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'})

    file = request.files['file']
    if not file or file.filename == '':
        return jsonify({'error': 'No selected file'})

    try:
        img = Image.open(file).convert('RGB')
        img_tensor = transform(img).unsqueeze(0).to(device)

        # 预测结果
        with torch.no_grad():
            output = model(img_tensor)
            prob = torch.softmax(output, dim=1)
            pred = torch.argmax(prob).item()
            confidence = prob[0, pred].item()

        # 生成Grad-CAM
        model.train()
        grad_cam, _ = generate_grad_cam(model, img_tensor)
        model.eval()

        # 使用OpenCV保存图像，避免Matplotlib多线程问题
        img_array = np.array(img.resize((224, 224))) / 255.0
        heatmap = cv2.applyColorMap(np.uint8(255 * grad_cam), cv2.COLORMAP_JET)
        cam = cv2.addWeighted(heatmap, 0.5, (img_array * 255).astype(np.uint8), 0.5, 0)
        cv2.imwrite('static/gradient.jpg', cam)

        return jsonify({
            'class': ID2CN[pred],
            'confidence': float(confidence),
            'grad_cam': 'static/gradient.jpg'
        })

    except Exception as e:
        return jsonify({'error': str(e)})

# 清理资源
@app.teardown_request
def cleanup(_):
    plt.close('all')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, threaded=True, use_reloader=False)